# https://juejin.cn/post/7178664815559180344
import torch
from torch import nn, optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
# 加载图片
from PIL import Image
from torchvision import models

"""
一、导入手写数据集
"""
train_dataset = datasets.MNIST(root='../data',
                               train=True,
                               transform=transforms.ToTensor(),
                               download=False)
test_dataset = datasets.MNIST(root='../data',
                              train=False,
                              transform=transforms.ToTensor(),
                              download=False)

"""
查看数据是否下载成功：
"""
# # 检查数据集大小
# print(f"训练集样本数: {len(train_dataset)}")  # 应输出 60000
# print(f"测试集样本数: {len(test_dataset)}")  # 应输出 10000
# # 检查单个样本的格式
# image, label = train_dataset[0]
# print(f"图像形状: {image.size}")  # 应输出 (28, 28)（未转Tensor时是PIL格式）
# print(f"标签: {label}")  # 应输出 0~9 的数字

"""
二、定义数据包装器
"""
batch_size = 64
train_loader = DataLoader(dataset=train_dataset,
                          batch_size=batch_size,
                          shuffle=True)
test_loader = DataLoader(dataset=test_dataset,
                         batch_size=batch_size,
                         shuffle=False)
"""
三、查看数据维度
"""
# for i, data in enumerate(train_loader):
#     inputs, labels = data
#     print(f"第 {i + 1} 批次的输入形状: {inputs.shape}")  # 应输出 (64, 1, 28, 28)
#     print(f"第 {i + 1} 批次的标签形状: {labels.shape}")  # 应输出 (64,)

"""
四、定义模型:卷积网络层
"""


class CNN(nn.Module):
    # 初始化相关网络层
    def __init__(self):
        # 继承父类初始化方法
        super(CNN, self).__init__()
        self.conv1 = nn.Sequential(
            nn.Conv2d(1, 64, 5, 1, 2),
            nn.ReLU(),
            nn.MaxPool2d(2, 2)
        )
        self.conv2 = nn.Sequential(
            nn.Conv2d(64, 128, 5, 1, 2),
            nn.ReLU(),
            nn.MaxPool2d(2, 2)
        )
        self.fc1 = nn.Sequential(
            nn.Linear(128 * 7 * 7, 1000),
            nn.ReLU(),
            nn.Dropout(p=0.2)
        )
        self.fc2 = nn.Sequential(
            nn.Linear(1000, 10),
            nn.Softmax(dim=1)
        )

    def forward(self, x):
        x = self.conv1(x)
        x = self.conv2(x)
        x = x.reshape(x.shape[0], -1)
        x = self.fc1(x)
        x = self.fc2(x)
        return x


"""
五、定义模型、损失函数和优化器
"""
# 学习率
LR = 0.0003
model = CNN()
loss_fn = nn.CrossEntropyLoss()  # 损失函数: 交叉熵损失函数
optimizer = optim.Adam(model.parameters(), lr=LR)

"""
六、训练、测试函数
"""


def model_train():
    model.train()
    for i, data in enumerate(train_loader):
        inputs, targets = data
        outputs = model(inputs)
        loss = loss_fn(outputs, targets)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()


def model_test():
    # model.eval() 会关闭 Dropout 和 BatchNorm 的随机性，但不自动关闭梯度计算。
    model.eval()
    # 测试训练集准确率（可选）
    correct = 0
    with torch.no_grad():  # 关键修改！禁止梯度计算
        for i, data in enumerate(train_loader):
            inputs, targets = data
            outputs = model(inputs)
            _, predicted = torch.max(outputs.data, 1)
            correct += (predicted == targets).sum()
        print(f"准确率：{100. * correct.item() / len(train_dataset)}%")
    # 测试测试集准确率
    correct = 0
    with torch.no_grad():  # 关键修改！禁止梯度计算
        for i, data in enumerate(test_loader):
            inputs, targets = data
            outputs = model(inputs)
            # 参数 1：表示沿着第1个维度（即“类别维度”）求最大值。
            """
            第一个值（用 _ 接收）：最大值的具体数值（这里不需要，所以用占位符 _ 忽略）。
            第二个值 predicted：最大值所在的索引（即模型预测的数字）。
            """
            _, predicted = torch.max(outputs.data, 1)
            # correct 累加 → 最终得到总正确数，用于计算准确率。
            correct += (predicted == targets).sum()
        print(f"准确率：{100. * correct.item() / len(test_dataset)}%")


"""
七、网络迭代训练
"""
# # 训练10次模型，不断修改权重，使得模型效果更优
# for epoch in range(10):
#     print(f"第{epoch + 1}次训练开始...")
#     model_train()
#     model_test()

"""
保存模型
"""
# 设定模型保存位置
model_save_path = './model_save/mnist_cnn_net01.pth'
# torch.save(model.state_dict(), model_save_path)
# print(f"模型保存成功，保存位置：{model_save_path}")

"""
九、加载模型
"""
# # 重新定义模型，加载保存的参数
# 因为是在同一个文件中，所以可以直接导入
# model = CNN()
# loss_fn = nn.CrossEntropyLoss()
# optimizer = optim.Adam(model.parameters(), lr=0.001)
model.load_state_dict(torch.load(model_save_path, weights_only=True))

"""
模型测试
"""
# 将图片转换为tensor
transform = transforms.Compose([transforms.ToTensor()])
# 打开图片
img = Image.open("../images/5.png")
# 改变图片大小
img = img.resize((28, 28))
# 图片灰度化
img = img.convert("L")
# 将图片增加一个维度
img = transform(img).unsqueeze(0)
print(img.shape)
out = model(img)
_, pred = torch.max(out, dim=1)
print(f"该数字是：{pred.item()}")
